package com.shujia.train

import com.alibaba.alink.common.MLEnvironmentFactory
import com.alibaba.alink.operator.batch.BatchOperator
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp
import com.alibaba.alink.operator.batch.evaluation.EvalBinaryClassBatchOp
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp
import com.alibaba.alink.operator.common.evaluation.BinaryClassMetrics
import com.alibaba.alink.pipeline.classification.NaiveBayesTextClassifier
import com.alibaba.alink.pipeline.nlp.{DocCountVectorizer, Segment, StopWordsRemover}
import com.alibaba.alink.pipeline.{Pipeline, PipelineModel}
import org.apache.flink.api.java.ExecutionEnvironment

object TrainModel {
  def main(args: Array[String]): Unit = {

    /**
      * 构建alink环境
      *
      */

    val env: ExecutionEnvironment = MLEnvironmentFactory.getDefault.getExecutionEnvironment

    /**
      * 读取数据
      *
      */

    val schema: String = "label Double,text String"

    /**
      * CsvSourceBatchOp 相当于flink的  DataSet
      *
      */
    val data: CsvSourceBatchOp = new CsvSourceBatchOp()
      .setFilePath("data/train.txt") //数据路径
      .setSchemaStr(schema) //数据列的描述
      .setFieldDelimiter("\t") //列分割方式


    //将数据切分成训练集和测试机
    val splitBatchOp: SplitBatchOp = new SplitBatchOp().setFraction(0.8)
    splitBatchOp.linkFrom(data)

    //获取测试集
    val testData: BatchOperator[_] = splitBatchOp.getSideOutput(0)

    System.out.println("训练集数据量：" + splitBatchOp.count)
    System.out.println("测试集数据量：" + testData.count)




    /**
      * Pipeline  流水线  --->  机器学习的流水线
      *
      * Segment  中文分词器
      * StopWordsRemover  去除停留词
      * DocCountVectorizer 将数据转换成向量
      * NaiveBayesTextClassifier 贝叶斯分类算法
      */


    val pipline: Pipeline = new Pipeline()
      .add(new Segment().setSelectedCol("text"))
      .add(new StopWordsRemover().setSelectedCol("text"))
      .add(new DocCountVectorizer().setSelectedCol("text"))
      .add(new NaiveBayesTextClassifier()
        .setLabelCol("label") //标签列
        .setVectorCol("text") //特征列
        .setPredictionCol("prediction") //预测列
        .setPredictionDetailCol("predictionDetail") //预测详细列
      )


    //将训练集带入流水线训练模型,  底层使用的是flink的离线计算
    val model: PipelineModel = pipline.fit(splitBatchOp)

    /**
      * 评估模型的准确率
      *
      */

    val textOp: BatchOperator[_ <: BatchOperator[_]] = model.transform(testData)

    /**
      * EvalBinaryClassBatchOp  用于判断二分类准确率的工具
      *
      */
    //模型评估
    val metrics: BinaryClassMetrics = new EvalBinaryClassBatchOp()
      .setLabelCol("label")
      .setPredictionDetailCol("predictionDetail")
      .linkFrom(textOp)
      .collectMetrics()


    println("AUC:" + metrics.getAuc)
    println("KS:" + metrics.getKs)
    println("PRC:" + metrics.getPrc)
    println("Accuracy:" + metrics.getAccuracy) //模型准确率
    println("Macro Precision:" + metrics.getMacroPrecision)
    println("Micro Recall:" + metrics.getMicroRecall)
    println("Weighted Sensitivity:" + metrics.getWeightedSensitivity)

    /**
      * 保存模型
      *
      */

    model.save("data/model")

    //启动任务
    BatchOperator.execute()


  }
}
